/*
 * simple MPI File distributor
 *
 * copy a directory to all other nodes
 * call with
 *
 *	mpirun -np <nodes> mpifilepush source target
 *
 * where source is a path to the files, and target is the directory
 * where the files should be put to. rsync like syntax applies, so
 * </some/dir/> will copy all files below </some/dir/> while
 * </some/dir> will copy the directory <dir> as well.
 *
 * Example:
 *    src = /data/testfiles/
 *    target = /tmp/testfiles
 * and
 *    src = /data/testfiles
 *    target = /tmp
 * copy both all files from within /data/testfiles into /tmp/testfiles
 * on the other nodes.
 *
 * so you propably want something like this with intel MPI
 *
 * cp -r /data/input /scratch/rundir    # nice for cache warmup as well
 * mpirun /scratch/rundir /scratch
 *
 * to have the file on all nodes! (as local node does not receive the data)
 *
 * the code can run several times on one host without doing any harm,
 * it will copy only once per node.
 *
 * (c) holger berger 2013
 *
 */

#include "mpi.h"

#include <stdio.h>
#include <stdlib.h>
#include <string>
#include <string.h>

#include <map>
#include <set>

using namespace std;

const int buffersize = (4*1024*1024);

MPI_Comm host_comm;
MPI_Group host_group;

/* helper to send a string */
void MPI_SendString(int dest, char *string)
{
	int len;
	len=strlen(string)+1;
	MPI_Send(&len, 1, MPI_INT, dest, 0, MPI_COMM_WORLD);
	MPI_Send(string, len, MPI_CHAR, dest, 0, MPI_COMM_WORLD);

}

/* helper to receive a string, memory has to be freed external! */
void MPI_RecvString(int src, char **string)
{
	int len;
	MPI_Status status;

	MPI_Recv(&len, 1, MPI_INT, src, 0, MPI_COMM_WORLD, &status);
	*string=(char*)malloc(len);
	MPI_Recv(*string, len, MPI_CHAR, src, 0, MPI_COMM_WORLD, &status);
}

/* gather a table of all hostnames on rank 0 */
void Init_host_group()
{
	int count,rank,len,i,j,nodes;
	int *rank_array;
	unsigned long hostid=0;
	unsigned long *hostid_array, *node_array;
	char hostname[MPI_MAX_PROCESSOR_NAME];
	MPI_Status status;

	MPI_Comm_rank(MPI_COMM_WORLD, &rank);
	MPI_Comm_size(MPI_COMM_WORLD, &count);

	MPI_Get_processor_name(hostname,&len);
	/* sdbm hash used because of its simplicity and good properties */
	for (i=0;i<len;i++) {
		hostid = (int) hostname[i] +
		  (hostid << 6) + (hostid << 16) - hostid;
	}

	hostid_array = (unsigned long *)malloc(sizeof(unsigned long) * count);
	MPI_Allgather(&hostid, 1, MPI_UNSIGNED_LONG, hostid_array,
		      1, MPI_UNSIGNED_LONG, MPI_COMM_WORLD);

	MPI_Group orig_group;
	MPI_Comm_group(MPI_COMM_WORLD, &orig_group);
	node_array = (unsigned long *)malloc(sizeof(unsigned long) * count);
	rank_array = (int *)malloc(sizeof(int) * count);
	nodes = 0;
	node_array[0]=hostid_array[0];
	for (i=0;i<count;i++) {
		for (j=0;j<nodes;j++) {
			if(hostid_array[i]==node_array[j])
				break;
		}
		if (j==nodes) {
			node_array[nodes] = hostid_array[i];
			rank_array[nodes] = i;
			nodes++;
		}
	}
	MPI_Group_incl(orig_group, nodes, rank_array, &host_group);
	MPI_Comm_create(MPI_COMM_WORLD, host_group, &host_comm);
}

int main(int argc, char **argv)
{
	int rank,count,host_rank;
	char* buf;
	char cmd[8192];
	FILE *fh;
	size_t len,total=0;
	double t1,t2,t3;

	buf = (char *)malloc(buffersize+sizeof(size_t));

	MPI_Init(&argc, &argv);

	MPI_Comm_rank(MPI_COMM_WORLD, &rank);
	MPI_Comm_size(MPI_COMM_WORLD, &count);

	Init_host_group();

	if(argc!=3 && rank==0) {
		printf("usage: %s src dst\n", argv[0]);
		printf("	copies files recursively from src to dst\n");
		printf("	where   src/dir dst \n");
		printf("	and     src/dir/ dst/dir \n");
		printf("	do the same\n");
		exit(0);
	}

	MPI_Barrier(MPI_COMM_WORLD);
	t1 = MPI_Wtime();
	MPI_Group_rank (host_group, &host_rank);
	if(rank==0) {
		if (host_rank == MPI_UNDEFINED)
			exit(1);
		if(argv[1][strlen(argv[1])-1]=='/') {
			snprintf(cmd, 8192, "(cd %s; tar cf - *)", argv[1]);
		} else {
			char copy[1024];
			strcpy(copy, argv[1]);
			*strrchr(copy,'/')=0;
			snprintf(cmd, 8192, "(cd %s; tar cf - %s)",
				 copy, strrchr(argv[1],'/')+1);
		}
		fh = popen(cmd,"r");
		if(fh<=0) {
			fprintf(stderr,"problems with pipe");
			MPI_Abort(MPI_COMM_WORLD, 1);
		}
		while(!feof(fh)) {
			len=fread(buf, (size_t)1, buffersize, fh);
			*(size_t*)(buf+buffersize)=len;
			MPI_Bcast(buf, buffersize+sizeof(size_t),
				  MPI_BYTE, 0, host_comm);
			total+=len;
		}
		if(len==buffersize) {
			*(size_t*)(buf+buffersize)=0;
			MPI_Bcast(buf, buffersize+sizeof(size_t),
				  MPI_BYTE, 0, host_comm);
		}
		pclose(fh);
	} else {
		if(host_rank != MPI_UNDEFINED) {
			snprintf(cmd, 8192, "(cd %s; tar xf -)", argv[2]);
			fh = popen(cmd,"w");
			while(1) {
				MPI_Bcast(buf, buffersize+sizeof(size_t),
					  MPI_BYTE, 0, host_comm);
				len = *(size_t*)(buf+buffersize);
				fwrite(buf, (size_t)1, len, fh);
				if(len<buffersize) break;
			}
			pclose(fh);
		}
	}
	t2 = MPI_Wtime();
	MPI_Barrier(MPI_COMM_WORLD);
	t3 = MPI_Wtime();

	int hostcount=0;
	MPI_Group_size( host_group, &hostcount) ;

	if(rank==0) {
		printf("transfered %ld bytes in %lf seconds, %lf MB/s\n", total,
		       t2-t1, ((double)total/(t2-t1))/(1024.0*1024.0));
		printf("aggregated BW was %lf MB/s\n",
		       (hostcount-1)*((double)total/(t3-t1))/(1024.0*1024.0));
	}

	MPI_Finalize();
}
